import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import json
from NER.config.config import Config
from NER.util.common import build_data
import torch
from torch.nn.utils.rnn import pad_sequence

config = Config()
datas, word2id = build_data()


class NerDataset(Dataset):
    def __init__(self, datas):
        super().__init__()
        self.datas = datas

    def __len__(self):
        return len(self.datas)

    def __getitem__(self, item):
        x = self.datas[item][0]
        y = self.datas[item][1]
        return x, y


def collate_fn(batch):
    # print(batch)
    # 需要将数据向量化
    x_train = [torch.tensor([word2id[char] for char in data[0]]) for data in batch]
    y_train = [torch.tensor([config.tag2id[label] for label in data[1]]) for data in batch]
    # print('x_train-->', x_train)
    # print("y_train-->", y_train)
    # 补齐input_ids, 使用0作为填充值
    input_ids_padded = pad_sequence(x_train, batch_first=True, padding_value=0)
    # # 补齐labels，使用0作为填充值
    labels_padded = pad_sequence(y_train, batch_first=True, padding_value=-100)
    # 创建attention mask
    attention_mask = (input_ids_padded != 0).long()
    return input_ids_padded, labels_padded, attention_mask


def getdataloader():
    train_dataset = NerDataset(datas=datas[:6200])
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=config.batch_size, shuffle=True,
                                  collate_fn=collate_fn,
                                  drop_last=True)

    dev_dataset = NerDataset(datas=datas[6200:])
    dev_dataloader = DataLoader(dataset=dev_dataset, batch_size=config.batch_size, shuffle=True, collate_fn=collate_fn,
                                drop_last=True)
    return train_dataloader, dev_dataloader


if __name__ == '__main__':
    getdataloader()
